package main

import (
	"fmt"
	"log"
	"os"

	"client2/zkp"

	"github.com/consensys/gnark-crypto/ecc"
	"github.com/consensys/gnark/backend/groth16"
	"github.com/consensys/gnark/frontend"
	"github.com/consensys/gnark/frontend/cs/r1cs"
)

// ⚠️ 警告: 此工具仅供开发团队使用！
// 客户端不应该运行此程序！
//
// 用途: 执行 Trusted Setup，生成全局密钥对
// 运行: go run ./cmd/setup_keys/main.go
// 输出:
//   - auth_proving.key   (分发给所有客户端)
//   - auth_verifying.key (部署到 Edge 服务端)

const (
	ProvingKeyFile   = "auth_proving.key"
	VerifyingKeyFile = "auth_verifying.key"
)

func main() {
	fmt.Println("============================================================")
	fmt.Println("⚠️  ZKP Trusted Setup 工具")
	fmt.Println("============================================================")
	fmt.Println()
	fmt.Println("此工具仅供开发团队使用！")
	fmt.Println("客户端不应该运行此程序！")
	fmt.Println()

	// 检查是否已存在密钥文件
	if fileExists(ProvingKeyFile) || fileExists(VerifyingKeyFile) {
		fmt.Println("⚠️  检测到现有密钥文件:")
		if fileExists(ProvingKeyFile) {
			fmt.Printf("   - %s\n", ProvingKeyFile)
		}
		if fileExists(VerifyingKeyFile) {
			fmt.Printf("   - %s\n", VerifyingKeyFile)
		}
		fmt.Println()
		fmt.Print("是否删除并重新生成? (yes/no): ")
		var response string
		fmt.Scanln(&response)
		if response != "yes" {
			fmt.Println("已取消")
			return
		}

		// 删除旧密钥
		os.Remove(ProvingKeyFile)
		os.Remove(VerifyingKeyFile)
		fmt.Println("✅ 已删除旧密钥文件")
		fmt.Println()
	}

	// 执行 Trusted Setup
	fmt.Println("开始 Trusted Setup...")
	fmt.Println("这可能需要 2-5 分钟，请耐心等待...")
	fmt.Println()

	pk, vk, err := setupKeys()
	if err != nil {
		log.Fatalf("❌ Trusted Setup 失败: %v", err)
	}

	// 保存密钥
	fmt.Println("保存密钥文件...")
	if err := saveProvingKey(pk, ProvingKeyFile); err != nil {
		log.Fatalf("❌ 保存 proving key 失败: %v", err)
	}
	fmt.Printf("✅ %s 保存成功\n", ProvingKeyFile)

	if err := saveVerifyingKey(vk, VerifyingKeyFile); err != nil {
		log.Fatalf("❌ 保存 verifying key 失败: %v", err)
	}
	fmt.Printf("✅ %s 保存成功\n", VerifyingKeyFile)

	// 显示文件信息
	fmt.Println()
	fmt.Println("============================================================")
	fmt.Println("✅ Trusted Setup 完成！")
	fmt.Println("============================================================")
	fmt.Println()
	fmt.Println("生成的文件:")

	if info, err := os.Stat(ProvingKeyFile); err == nil {
		fmt.Printf("  📄 %s (%d bytes)\n", ProvingKeyFile, info.Size())
		fmt.Println("     → 分发给所有客户端")
	}

	if info, err := os.Stat(VerifyingKeyFile); err == nil {
		fmt.Printf("  📄 %s (%d bytes)\n", VerifyingKeyFile, info.Size())
		fmt.Println("     → 部署到 Edge 服务端")
	}

	fmt.Println()
	fmt.Println("下一步:")
	fmt.Println("  1. 将 auth_proving.key 分发给所有客户端")
	fmt.Println("  2. 将 auth_verifying.key 部署到 Edge 服务端")
	fmt.Println("  3. 妥善保管 Trusted Setup 环境（或销毁）")
	fmt.Println()
}

// setupKeys 执行 Trusted Setup，生成密钥对
func setupKeys() (groth16.ProvingKey, groth16.VerifyingKey, error) {
	// 创建空电路
	circuit := &zkp.AuthCircuit{}

	// 编译电路为 R1CS
	fmt.Println("📋 编译电路...")
	ccs, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, circuit)
	if err != nil {
		return nil, nil, fmt.Errorf("编译电路失败: %v", err)
	}
	fmt.Printf("✅ 电路编译成功 (约束数: %d)\n", ccs.GetNbConstraints())

	// 执行可信设置 (Trusted Setup)
	fmt.Println("🔐 执行 Trusted Setup (Groth16)...")
	pk, vk, err := groth16.Setup(ccs)
	if err != nil {
		return nil, nil, fmt.Errorf("Trusted Setup 失败: %v", err)
	}
	fmt.Println("✅ Trusted Setup 完成")

	return pk, vk, nil
}

// 辅助函数

func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	return err == nil
}

func saveProvingKey(pk groth16.ProvingKey, filename string) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = pk.WriteTo(file)
	return err
}

func saveVerifyingKey(vk groth16.VerifyingKey, filename string) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = vk.WriteTo(file)
	return err
}
